import tkinter as tk
import tkinter.filedialog
import numpy as np
import torch
import time
from models import SRCNN
import PIL.Image as Image
import matplotlib.pyplot as plt
import torch.backends.cudnn as cudnn

def convert_rgb_to_ycbcr(img):
    if type(img) == np.ndarray:
        y = 16. + (64.738 * img[:, :, 0] + 129.057 * img[:, :, 1] + 25.064 * img[:, :, 2]) / 256.
        cb = 128. + (-37.945 * img[:, :, 0] - 74.494 * img[:, :, 1] + 112.439 * img[:, :, 2]) / 256.
        cr = 128. + (112.439 * img[:, :, 0] - 94.154 * img[:, :, 1] - 18.285 * img[:, :, 2]) / 256.
        return np.array([y, cb, cr]).transpose([1, 2, 0])
    elif type(img) == torch.Tensor:
        if len(img.shape) == 4:
            img = img.squeeze(0)
        y = 16. + (64.738 * img[0, :, :] + 129.057 * img[1, :, :] + 25.064 * img[2, :, :]) / 256.
        cb = 128. + (-37.945 * img[0, :, :] - 74.494 * img[1, :, :] + 112.439 * img[2, :, :]) / 256.
        cr = 128. + (112.439 * img[0, :, :] - 94.154 * img[1, :, :] - 18.285 * img[2, :, :]) / 256.
        return torch.cat([y, cb, cr], 0).permute(1, 2, 0)
    else:
        raise Exception('Unknown Type', type(img))

def convert_ycbcr_to_rgb(img):
    if type(img) == np.ndarray:
        r = 298.082 * img[:, :, 0] / 256. + 408.583 * img[:, :, 2] / 256. - 222.921
        g = 298.082 * img[:, :, 0] / 256. - 100.291 * img[:, :, 1] / 256. - 208.120 * img[:, :, 2] / 256. + 135.576
        b = 298.082 * img[:, :, 0] / 256. + 516.412 * img[:, :, 1] / 256. - 276.836
        return np.array([r, g, b]).transpose([1, 2, 0])
    elif type(img) == torch.Tensor:
        if len(img.shape) == 4:
            img = img.squeeze(0)
        r = 298.082 * img[0, :, :] / 256. + 408.583 * img[2, :, :] / 256. - 222.921
        g = 298.082 * img[0, :, :] / 256. - 100.291 * img[1, :, :] / 256. - 208.120 * img[2, :, :] / 256. + 135.576
        b = 298.082 * img[0, :, :] / 256. + 516.412 * img[1, :, :] / 256. - 276.836
        return torch.cat([r, g, b], 0).permute(1, 2, 0)
    else:
        raise Exception('Unknown Type', type(img))

def choosepic():
    global image
    path_ = tkinter.filedialog.askopenfilename()
    path.set(path_)
    image = Image.open(entry.get()).convert('RGB')
    print('original image size:', image.size)
    plt.figure("original image"), plt.imshow(image), plt.axis('off')
    plt.show()

def SRpic(scale):
    global output
    global image_bicubic
    image_bicubic = image.resize((image.width * scale, image.height * scale), resample=Image.BICUBIC)

    cudnn.benchmark = True
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    model = SRCNN().to(device)
    state_dict = model.state_dict()
    for n, p in torch.load('srcnn_x4.pth', map_location=lambda storage, loc: storage).items():
        if n in state_dict.keys():
            state_dict[n].copy_(p)
        else:
            raise KeyError(n)
    model.eval()
    print('model load finished')

    image_b = np.array(image_bicubic).astype(np.float32)        # float32 [0,255] [H,W,C]
    ycbcr = convert_rgb_to_ycbcr(image_b)
    y = ycbcr[..., 0]
    y /= 255.
    y = torch.from_numpy(y).to(device)
    y = y.unsqueeze(0).unsqueeze(0)

    t0 = time.time()
    with torch.no_grad():
        preds = model(y).clamp(0.0, 1.0)
    t1 = time.time()
    print('Processing Time ={:.3f}s'.format(t1 - t0))
    preds = preds.mul(255.0).cpu().numpy().squeeze(0).squeeze(0)
    output = np.array([preds, ycbcr[..., 1], ycbcr[..., 2]]).transpose([1, 2, 0])
    output = np.clip(convert_ycbcr_to_rgb(output), 0.0, 255.0).astype(np.uint8)
    output = Image.fromarray(output)

    print('bicubic image size',image_bicubic.size)
    plt.figure("bicubic image"), plt.imshow(image_bicubic), plt.axis('off')
    plt.show()
    print('sr image size',output.size)
    plt.figure("sr image"), plt.imshow(output), plt.axis('off')
    plt.show()

def savepic(imgtype):
    fname = tkinter.filedialog.asksaveasfilename(title=u'保存文件', filetypes=[("PNG", ".png")])
    if imgtype=='Bicubic':
        image_bicubic.save(fname + '.png', 'PNG')
    elif imgtype=='SRCNN':
        output.save(fname + '.png', 'PNG')

if __name__ == '__main__':
    app = tk.Tk()                           # 生成tk界面 app即主窗口
    app.title("SRCNN_GUI")                  # 修改窗口titile
    app.geometry("400x300")                 # 设置主窗口的大小和位置

    path = tk.StringVar()                   # 显示简单的文本（读取图像的路径）
    entry = tk.Entry(app, state='readonly', text=path,width = 100)
    entry.pack()

    # 设置三个按钮
    b0 = tk.Button(app, text='选择图片', command=choosepic)
    b0.pack()

    b1 = tk.Button(app, text='展示超分后图片', fg='black', command=lambda: SRpic(scale=4))
    b1.place(x=60, y=135, width=100, height=30)
    b1.pack()

    b2 = tk.Button(app, text='保存插值结果图片', fg='black', command=lambda: savepic('Bicubic'))
    b2.place(x=100, y=135, width=100, height=30)
    b2.pack()

    b3 = tk.Button(app, text='保存超分结果图片', fg='black', command=lambda: savepic('SRCNN'))
    b3.place(x=140, y=135, width=100, height=30)
    b3.pack()

    app.mainloop()                          # Call the mainloop of Tk.